# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import csv
import datetime
import pathlib
import sys
from collections import namedtuple

from absl import (flags, app)

from coin.exchange.binance.kr_rest.product import BinanceProduct
from coin.exchange.bitflyer_v1.kr_rest.futures_product import BitflyerFuturesProduct
from coin.exchange.bitfinex_v2.kr_rest.product import BitfinexProduct
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.gdax.kr_rest.product import GdaxProduct
from coin.exchange.okex.kr_rest.product import OkexProduct
from coin.exchange.okex_futures.kr_rest.futures_product import OkexFuturesProduct
from coin.support.feed_tool.feed_stats.logic.util import *
from coin.strategy.mm.subscription import get_subscription_key
from coin.strategy.mm.dumper_base import run_from_archive

FLAGS = flags.FLAGS

CsvWriter = namedtuple('CsvWriter', ['writer', 'file', 'product'])


def get_ticksize(product):
  if isinstance(product, BitmexFuturesProduct):
    return 1
  elif isinstance(product, OkexFuturesProduct):
    return 0.5
  elif isinstance(product, BitflyerFuturesProduct):
    return 50
  else:
    raise ValueError()


class DumpStrategy(object):
  def __init__(self, products, csv_root, book_levels):
    self._products = products
    self._csv_root = csv_root
    self._book_levels = book_levels
    pathlib.Path(csv_root).mkdir(parents=True, exist_ok=True)
    self._csv_writers = {}

  def prepare_csv_file(self, product):
    csv_file = open('%s/%s.csv' % (self._csv_root, product.full_symbol), 'w', encoding='utf-8')
    writer = csv.writer(csv_file, dialect='excel')
    csv_header = ['timestamp']
    ticksize = get_ticksize(product)
    for level in range(self._book_levels):
      csv_header += [
          'bid%d_%f' % (level, ticksize),
          'bid_qty%d_%f' % (level, ticksize),
          'ask%d_%f' % (level, ticksize),
          'ask_qty%d_%f' % (level, ticksize)
      ]
    writer.writerow(csv_header)
    csv_file.flush()
    return writer, csv_file

  def on_book_reset(self, book_builder_name, book_builder):
    for product in self._products:
      if book_builder_name == get_subscription_key(product):
        book_builder.subscribe_aggregated_book(product, [get_ticksize(product)], self.on_book)
        exchange, symbol = to_book_builder_exchange_symbol(product)
        writer, csv_file = self.prepare_csv_file(product)
        self._csv_writers[(exchange, symbol)] = CsvWriter(writer, csv_file, product)

  def on_book(self, book):
    writer = self._csv_writers.get((book.exchange, book.symbol), None)
    assert writer is not None, (book.exchange, book.symbol)

    try:
      bid0 = book.bid0()
      ask0 = book.ask0()
      csv_line = [book.timestamp]
    except IndexError:
      return

    ticksize = get_ticksize(writer.product)
    bids = [[None, None]] * self._book_levels
    level = 0
    for key, value in book.get_aggregated_book(ticksize).iteritems_bid():
      assert key == value.price
      assert value.qty > 0
      bids[level] = [value.price, value.qty]
      level += 1
      if level >= self._book_levels:
        break
    asks = [[None, None]] * self._book_levels
    level = 0
    for key, value in book.get_aggregated_book(ticksize).iteritems_ask():
      assert key == value.price
      assert value.qty > 0
      asks[level] = [value.price, value.qty]
      level += 1
      if level >= self._book_levels:
        break

    for level in range(self._book_levels):
      csv_line += [bids[level][0], bids[level][1], asks[level][0], asks[level][1]]

    writer.writer.writerow(csv_line)
    writer.file.flush()

  def close_csvs(self):
    for writer in self._csv_writers.values():
      writer.file.flush()
      writer.file.close()


def launch(trading_date, machine, csv_root, book_levels, hours):
  bitmex_product = BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')

  okex_futures_product = OkexFuturesProduct.FromStr('BTC-USD.QUARTER', trading_date)

  bitflyer_product = BitflyerFuturesProduct.FromStr("BTC-JPY.PERPETUAL")

  products = [bitmex_product, okex_futures_product, bitflyer_product]

  start_time = trading_date  # UTC 0 hour

  strategy = DumpStrategy(products, csv_root, book_levels)

  run_from_archive(products,
                   strategy.on_book_reset,
                   start_time,
                   start_time + datetime.timedelta(hours=hours),
                   machine=[machine])

  strategy.close_csvs()
  return 0


def main(argv):
  trading_date = FLAGS.trading_date
  machine = FLAGS.machine
  csv_root = FLAGS.csv_root
  book_levels = FLAGS.book_levels
  assert trading_date, '--trading_date must be specified.'
  assert machine, '--machine must be specified.'
  assert csv_root, '--csv_root must be specified.'
  print('Running for %s %s ...' % (trading_date, machine))
  sys.stdout.flush()
  trading_date = datetime.datetime.strptime(trading_date, '%Y%m%d')
  launch(trading_date, machine, csv_root, book_levels, FLAGS.hours)


if __name__ == '__main__':
  flags.DEFINE_string('trading_date', None, 'Trading date in form of %Y%m%d.')

  flags.DEFINE_string('machine', None, 'Instance machine name.')

  flags.DEFINE_string('csv_root', None, 'Output csv files root directory.')

  flags.DEFINE_integer('book_levels', 5, 'How many levels of book to print in csv.')

  flags.DEFINE_integer('hours', 24, 'Time span in hours from 00:00:00 UTC.')

  app.run(main)
